package org.os.framework.crypto.algorithm;

import lombok.SneakyThrows;
import org.os.framework.crypto.model.Crypto;
import org.os.framework.crypto.model.Signature;
import org.os.framework.crypto.model.AbstractStatic;

import cn.hutool.core.lang.Assert;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA
 *
 * @author 魏祖潇
 */
public class Rsa extends AbstractStatic implements Crypto<String>, Signature {

    /**
     * 构造器
     */
    public static class Builder {
        private String privateKey;
        private String publicKey;

        public Builder() {
        }

        public Builder privateKey(String privateKey) {
            this.privateKey = privateKey;
            return this;
        }

        public Builder publicKey(String publicKey) {
            this.publicKey = publicKey;
            return this;
        }

        public Rsa build() {
            Rsa rsa = new Rsa();
            Assert.isTrue(null != this.privateKey || null != this.publicKey, "至少输入一种密钥");
            if (null != this.privateKey) {
                rsa.privateKey(this.privateKey);
            }
            if (null != this.publicKey) {
                rsa.publicKey(this.publicKey);
            }
            return rsa;
        }

    }

    public static Builder builder() {
        return new Builder();
    }

    private Rsa() {
    }

    private final String algorithm = "RSA";

    /**
     * 私钥
     */
    private RSAPrivateKey privateKey;

    /**
     * 公钥
     */
    private RSAPublicKey publicKey;

    public Rsa(String privateKey, String publicKey) {
        if (null != privateKey) {
            privateKey(privateKey);
        }
        if (null != publicKey) {
            publicKey(publicKey);
        }
    }

    /**
     * 设置公钥
     * 
     * @param publicKey
     */
    @SneakyThrows
    private void publicKey(String publicKey) {
        Assert.notBlank(publicKey, "请输入公钥");
        if (isPemFormat(publicKey)) {
            publicKey = removePemHeadersAndDecode(publicKey);
        }
        byte[] publicKeyBytes = Base64.getDecoder().decode(publicKey);
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(publicKeyBytes);
        KeyFactory factory = KeyFactory.getInstance(algorithm);
        this.publicKey = (RSAPublicKey) factory.generatePublic(x509KeySpec);
    }

    /**
     * 设置私钥
     * 
     * @param privateKey
     */
    @SneakyThrows
    private void privateKey(String privateKey) {
        Assert.notBlank(privateKey, "请输入私钥");
        if (isPemFormat(privateKey)) {
            privateKey = removePemHeadersAndDecode(privateKey);
        }
        // 解码密钥字符串
        byte[] privateKeyBytes = Base64.getDecoder().decode(privateKey);
        // 创建密钥规格
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
        // 获取密钥工厂实例
        KeyFactory factory = KeyFactory.getInstance(algorithm);
        // 生成公钥和私钥对象
        this.privateKey = (RSAPrivateKey) factory.generatePrivate(pkcs8KeySpec);
    }

    @SneakyThrows
    @Override
    public String encrypt(String original) {
        Cipher cipher = Cipher.getInstance(algorithm);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        int length = original.getBytes().length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        int i = 0;
        byte[] cache;
        while (length - offSet > 0) {
            if (length - offSet > 117) {
                cache = cipher.doFinal(original.getBytes(), offSet, 117);
            } else {
                cache = cipher.doFinal(original.getBytes(), offSet, length - offSet);
            }
            out.write(cache, 0, cache.length);
            i++;
            offSet = i * 117;
        }
        byte[] bytes = out.toByteArray();
        out.close();
        return Base64.getEncoder().encodeToString(bytes);
    }

    @SneakyThrows
    @Override
    public String decrypt(String encrypted) {
        byte[] bytes = Base64.getDecoder().decode(encrypted);
        Cipher cipher = Cipher.getInstance(algorithm);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        int length = bytes.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段解密
        while (length - offSet > 0) {
            if (length - offSet > 128) {
                cache = cipher.doFinal(bytes, offSet, 128);
            } else {
                cache = cipher.doFinal(bytes, offSet, length - offSet);
            }
            out.write(cache, 0, cache.length);
            i++;
            offSet = i * 128;
        }
        byte[] s = out.toByteArray();
        out.close();
        return new String(s);
    }

    @SneakyThrows
    @Override
    public String sign(String original) {
        java.security.Signature s = java.security.Signature.getInstance(algorithm);
        s.initSign(privateKey);
        s.update(original.getBytes());
        return Base64.getEncoder().encodeToString(s.sign());
    }

    @SneakyThrows
    @Override
    public boolean verify(String original, String signature) {
        java.security.Signature s = java.security.Signature.getInstance(algorithm);
        s.initVerify(publicKey);
        s.update(original.getBytes());
        return s.verify(Base64.getDecoder().decode(signature));
    }

    public BigInteger encrypt(BigInteger m) {
        BigInteger e = publicKey.getPublicExponent();
        BigInteger n = publicKey.getModulus();
        return m.modPow(e, n);
    }

    public BigInteger decrypt(BigInteger c) {
        BigInteger d = privateKey.getPrivateExponent();
        BigInteger n = privateKey.getModulus();
        return c.modPow(d, n);
    }

    public BigInteger multiply(BigInteger a, BigInteger b) {
        return a.multiply(b);
    }

    /**
     * 辅助方法，用于处理PEM格式字符串
     * @param keyStr
     * @return
     */
    boolean isPemFormat(String keyStr) {
        return keyStr.contains("-----BEGIN RSA PUBLIC KEY-----") || keyStr.contains("-----BEGIN RSA PRIVATE KEY-----");
    }

    String removePemHeadersAndDecode(String pemStr) {
        String trimmed = pemStr.replaceAll("-----BEGIN (.*)-----\\s+|\\s+-----END (.*)-----(\\s+|\r\n)", "");
        return new String(Base64.getDecoder().decode(trimmed), StandardCharsets.UTF_8);
    }
}
